import tensorflow as tf
from tensorflow.keras.layers import Input, Conv2D, BatchNormalization, Activation, MaxPooling2D, Add, Flatten, Dense
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt

# 指定使用的 GPU 设备
physical_devices = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_visible_devices(physical_devices[0], 'GPU')


# ResNet卷积块：两个3x3的卷积层和恒等映射
def resnet_block(inputs, filters, kernel_size, strides):
	x = Conv2D(filters, kernel_size=kernel_size, strides=strides, padding='same')(inputs)
	x = BatchNormalization()(x)
	x = Activation('relu')(x)
	
	x = Conv2D(filters, kernel_size=kernel_size, strides=1, padding='same')(x)
	x = BatchNormalization()(x)
	
	# 恒等映射
	if strides != 1:
		shortcut = Conv2D(filters, kernel_size=1, strides=strides, padding='same')(inputs)
		shortcut = BatchNormalization()(shortcut)
	else:
		shortcut = inputs
	
	x = Add()([x, shortcut])
	x = Activation('relu')(x)
	return x


def build_resnet(input_shape, num_classes):
	inputs = Input(shape=input_shape)
	
	x = Conv2D(64, kernel_size=7, strides=2, padding='same')(inputs)
	x = BatchNormalization()(x)
	x = Activation('relu')(x)
	x = MaxPooling2D(pool_size=(3, 3), strides=2, padding='same')(x)
	
	x = resnet_block(x, filters=64, kernel_size=3, strides=1)
	x = resnet_block(x, filters=64, kernel_size=3, strides=1)
	x = resnet_block(x, filters=64, kernel_size=3, strides=1)
	
	x = resnet_block(x, filters=128, kernel_size=3, strides=2)
	x = resnet_block(x, filters=128, kernel_size=3, strides=1)
	x = resnet_block(x, filters=128, kernel_size=3, strides=1)
	
	x = resnet_block(x, filters=256, kernel_size=3, strides=2)
	x = resnet_block(x, filters=256, kernel_size=3, strides=1)
	x = resnet_block(x, filters=256, kernel_size=3, strides=1)
	
	x = resnet_block(x, filters=512, kernel_size=3, strides=2)
	x = resnet_block(x, filters=512, kernel_size=3, strides=1)
	x = resnet_block(x, filters=512, kernel_size=3, strides=1)
	
	x = BatchNormalization()(x)
	x = Activation('relu')(x)
	x = Flatten()(x)
	x = Dense(512, activation='relu')(x)
	x = Dense(num_classes, activation='softmax')(x)
	
	model = Model(inputs=inputs, outputs=x)
	return model


IMSIZE = 224
num_classes = 200

train_datagen = ImageDataGenerator(
	rescale=1. / 255,
	rotation_range=20,
	width_shift_range=0.2,
	height_shift_range=0.2,
	shear_range=0.2,
	zoom_range=0.2,
	horizontal_flip=True,
	fill_mode='nearest'
)
train_generator = train_datagen.flow_from_directory(
    '../data/flower_learn_data/trains',
    target_size=(IMSIZE, IMSIZE),
    batch_size=32,
    class_mode='categorical'
)

validation_datagen = ImageDataGenerator(rescale=1. / 255)
validation_generator = validation_datagen.flow_from_directory(
    '../data/flower_learn_data/tests',
    target_size=(IMSIZE, IMSIZE),
    batch_size=32,
    class_mode='categorical'
)

model = build_resnet((IMSIZE, IMSIZE, 3), num_classes)

model.compile(optimizer=Adam(learning_rate=0.005), loss='categorical_crossentropy', metrics=['accuracy'])

early_stopping = tf.keras.callbacks.EarlyStopping(
	monitor='val_loss',
	patience=3,
	restore_best_weights=True
)

history = model.fit(
	train_generator,
	steps_per_epoch=train_generator.samples // train_generator.batch_size,
	epochs=100,
	validation_data=validation_generator,
	validation_steps=validation_generator.samples // validation_generator.batch_size,
	callbacks=[early_stopping]
)

# 评估模型
accuracy = model.evaluate(validation_generator)[1]
print("准确率:", accuracy)

# 绘制训练过程中的准确率和损失曲线
plt.plot(history.history['accuracy'])
plt.plot(history.history['val_accuracy'])
plt.title('Model Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend(['Train', 'Validation'], loc='upper left')
plt.show()

plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('Model Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend(['Train', 'Validation'], loc='upper left')
plt.show()